You write custom CUDA kernels to replace PyTorch operators for speedups.
Implement RMS-Gated Normalization: Given a 2D tensor x of shape [B, D], compute the per-row RMS r = sqrt(mean(x^2, dim=1)), then a scalar gate g = sigmoid(gamma * r + beta) per row, and output y = x * g. The CUDA kernel must perform a block-level reduction to obtain RMS for each row (one block per row), then apply the gate to all elements in that row. Use contiguous memory layout, dynamic shared memory for the reduction, and avoid multiple kernel launches. Provide a PyTorch reference module that computes the same operation using nn.Parameters for gamma and beta. Ensure numerical stability and accuracy matching torch reference within rtol=1e-3.
